from dataclasses import dataclass
from threading import Thread
from typing import NamedTuple, Callable
from time import sleep
from copy import deepcopy

SEC_PER_TICK = 1

@dataclass(init=False)
class CpuTask:
    number: int
    cost_ticks: int
    remaining_ticks: int
    priority: int = 0 # 【优先队列】优先级越大越高。
    
    def __init__(self, number: int, cost_ticks: int, priority: int = 0):
        self.number = number
        self.cost_ticks = cost_ticks
        self.remaining_ticks = cost_ticks
        self.priority = priority

class FIFOScheduler:
    def __init__(
        self,
        tasks: list[CpuTask]|None = None,
        tick_duration: float = SEC_PER_TICK,
        on_tick: Callable[[int, CpuTask], None] = lambda a, b: None
    ):
        self.queue = deepcopy(tasks) if tasks else []
        self.current_task: CpuTask|None = None
        self.tick: int = 0
        self.on_tick = on_tick
        self.tick_duration = tick_duration
        self._stopping: bool = False
        self.__thread: Thread|None = None
        
    def append_task(self, task):
        self.queue.append(task)
        
    def start(self):
        self.__thread = Thread(target=self._run)
        self.__thread.start()
    
    def stop(self):
        if not self.__thread:
            return
        self._stopping = True
        self.__thread.join()
        self.__thread = None
        self.tick = 0
        self._stopping = False
    
    def clear(self):
        self.queue.clear()
        self.current_task = None
    
    def join(self):
        if self.__thread:
            self.__thread.join()
        
    def _run(self):
        while (self.queue or self.current_task) and not self._stopping:
            if not self.current_task:
                self.current_task = self.queue.pop(0)
            self.current_task.remaining_ticks -= 1
            self._debug_tick_print()
            self.on_tick(self.tick, self.current_task)
            if self.current_task.remaining_ticks == 0:
                self.current_task = None
                self.current_task = self.queue.pop(0) if self.queue else None
            self.tick += 1
            if self.tick_duration > 0:
                sleep(self.tick_duration)
        self._debug_tick_print()

    def _debug_tick_print(self):
        if self.current_task:
            print(f'FIFO: t={self.tick}, task={self.current_task.number}[{self.current_task.priority}], remaining={self.current_task.remaining_ticks}')
        else:
            print(f'FIFO: t={self.tick}, idle')

class PriorityScheduler(FIFOScheduler):
    def __init__(
        self,
        tasks: list[CpuTask]|None = None,
        tick_duration: float = SEC_PER_TICK,
        on_tick: Callable[[float, CpuTask], None] = lambda a, b: None,
    ):
        super().__init__(tasks, tick_duration, on_tick)
        self.queue.sort(key=lambda task: task.priority, reverse=True)
    
    def append_task(self, task):
        self.queue.append(task)
        self.queue.sort(key=lambda task: task.priority, reverse=True)
    
    def _debug_tick_print(self):
        if self.current_task:
            print(f'Priority: t={self.tick}, task={self.current_task.number}, remaining={self.current_task.remaining_ticks}')
        else:
            print(f'Priority: t={self.tick}, idle')


class RoundRobinScheduler(FIFOScheduler):
    def __init__(
        self,
        tasks: list[CpuTask]|None = None,
        tick_duration: float = SEC_PER_TICK,
        on_tick: Callable[[float, CpuTask], None] = lambda a, b: None,
        quantum_tick: int = 5
    ):
        super().__init__(tasks, tick_duration, on_tick)
        self.quantum_tick = quantum_tick
        self._remaining_quantum = quantum_tick
    
    def _run(self):
        while (self.queue or self.current_task) and not self._stopping:
            if not self.current_task:
                self.current_task = self.queue.pop(0)
            self.current_task.remaining_ticks -= 1
            self._remaining_quantum -= 1
            self._debug_tick_print()
            self.on_tick(self.tick, self.current_task)
            # 任务完成判断
            if self.current_task.remaining_ticks <= 0:
                self.current_task = None
                self._remaining_quantum = self.quantum_tick
            # 时间片用尽判断
            if self._remaining_quantum == 0:
                if self.current_task:
                    self.queue.append(self.current_task)
                    self.current_task = None
                self._remaining_quantum = self.quantum_tick
            self.tick += 1
            if self.tick_duration > 0:
                sleep(self.tick_duration)
        self._debug_tick_print()
        
    def _debug_tick_print(self):
        if self.current_task:
            print(f'Round Robin: t={self.tick}, task={self.current_task.number}, remaining={self.current_task.remaining_ticks}')
        else:
            print(f'Round Robin: t={self.tick}, idle')


if __name__ == '__main__':
    # scheduler = FIFOScheduler([
    #     CpuTask(1, 3, 3),
    #     CpuTask(2, 2, 2),
    #     CpuTask(3, 1, 1),
    # ])
    # scheduler.start()
    # scheduler.join()
    # print('---------------')
    # scheduler = PriorityScheduler([
    #     CpuTask(1, 3, 3, 1),
    #     CpuTask(2, 2, 2, 2),
    #     CpuTask(3, 1, 1, 3),
    # ])
    # scheduler.start()
    # scheduler.join()
    # print('---------------')
    scheduler = RoundRobinScheduler([
        CpuTask(number=1, cost_ticks=6),
        CpuTask(number=2, cost_ticks=4),
        CpuTask(number=3, cost_ticks=2)
    ], quantum_tick=4)
    scheduler.start()